import {GetProp, theme, Transfer, TransferProps, Tree, TreeDataNode} from "antd";
import React from "react";


type TransferItem = GetProp<TransferProps, 'dataSource'>[number];

export interface TreeTransferProps extends TransferProps{
    dataSource: TreeDataNode[];
    targetKeys: TransferProps['targetKeys'];
    onChange: TransferProps['onChange'];
}

// Customize Table Transfer
const isChecked = (selectedKeys: React.Key[], eventKey: React.Key) =>
    selectedKeys.includes(eventKey);

const generateTree = (
    treeNodes: TreeDataNode[] = [],
    checkedKeys: TreeTransferProps['targetKeys'] = [],
): TreeDataNode[] =>
    treeNodes.map(({children, ...props}) => ({
        ...props,
        disabled: checkedKeys.includes(props.key as string),
        children: generateTree(children, checkedKeys),
    }));

const TreeTransfer: React.FC<TreeTransferProps> = ({
                                                       dataSource,
                                                       targetKeys = [],
                                                       ...restProps
                                                   }) => {
    const {token} = theme.useToken();

    const transferDataSource: TransferItem[] = [];

    function flatten(list: TreeDataNode[] = []) {
        list.forEach((item) => {
            transferDataSource.push(item as TransferItem);
            flatten(item.children);
        });
    }

    flatten(dataSource);

    return (
        <Transfer
            {...restProps}
            targetKeys={targetKeys}
            dataSource={transferDataSource}
            className="tree-transfer"
            render={(item) => item.title!}
            showSelectAll={true}
        >
            {({direction, onItemSelect, selectedKeys}) => {
                if (direction === 'left') {
                    const checkedKeys = [...selectedKeys, ...targetKeys];
                    return (
                        <div style={{padding: token.paddingXS}}>
                            <Tree
                                blockNode
                                checkable
                                checkStrictly
                                defaultExpandAll
                                checkedKeys={checkedKeys}
                                treeData={generateTree(dataSource, targetKeys)}
                                onCheck={(_, {node: {key}}) => {
                                    onItemSelect(key as string, !isChecked(checkedKeys, key));
                                }}
                                onSelect={(_, {node: {key}}) => {
                                    onItemSelect(key as string, !isChecked(checkedKeys, key));
                                }}

                            />
                        </div>
                    );
                }
            }}
        </Transfer>
    );
};


export default TreeTransfer
